package com.uyayo.framework.limiter.support;

import com.google.common.collect.Lists;
import com.uyayo.framework.limiter.exception.RequestLimiterException;
import org.redisson.api.RScript;
import org.redisson.api.RedissonClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.invoke.MethodHandles;

/**
 * uyayo-framework-limiter
 * 请求限制帮助类
 *
 * @author hfj
 *         date 2017/3/24
 */

public class RequestLimiterHelper {

    private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
    private static final String SCRIPT_LUA = "local current;local number =ARGV[1];" +
            "current = redis.call('incr',KEYS[1]);" +
            "if tonumber(current) == 1 then redis.call('expire',KEYS[1],number) end;" +
            "return current;";

    String SCRIPT_LUA_FOR_EVER = "local current;current = redis.call('incr',KEYS[1]);return current;";

    private static final String KEY_PREFIX = "request_limiter";
    private static final String COLON = ":";
    private final RedissonClient redissonClient;


    public RequestLimiterHelper(RedissonClient redissonClient) {
        this.redissonClient = redissonClient;
    }

    /**
     * 递增请求次数,超过最大允许次数会抛出RequestLimiterException 异常
     * 常用: 单位时间: 300s, 上限次数: 1
     *
     * @param identifier  限制标识
     * @param transaction 事务类型(可以根据需求,传入当前 URI)u
     */
    public void incrRequestLimiter(String transaction, String identifier) {
        this.incrRequestLimiter(transaction, identifier, 300, 1);
    }

    /**
     * 递增请求次数,超过最大允许次数会抛出RequestLimiterException 异常
     *
     * @param transaction 事务类型(可以根据需求,传入当前 URI)
     * @param identifier  限制标识
     * @param seconds     单位时间(s)
     * @param maxCount    单位时间内上限次数
     */
    public void incrRequestLimiter(String transaction, String identifier, long seconds, long maxCount) {
        final String key = KEY_PREFIX+transaction + COLON + identifier;
        this.incrRequestLimiter(key, seconds, maxCount);
    }

    /**
     * 根据key返回执行次数
     * @param key
     * @return
     */
    public Long incr(String key) {
        Long count = redissonClient.getScript()
                .eval(RScript.Mode.READ_WRITE, SCRIPT_LUA_FOR_EVER,
                        RScript.ReturnType.INTEGER, Lists.newArrayList(key));
        return count;
    }

    /**
     * 递增请求次数,超过最大允许次数会抛出RequestLimiterException 异常
     *
     * @param key      请求类型KEY
     * @param seconds  单位时间(s)
     * @param maxCount 单位时间内上限次数
     */
    public void incrRequestLimiter(String key, long seconds, long maxCount) {
        Long count = redissonClient.getScript()
                .eval(RScript.Mode.READ_WRITE, SCRIPT_LUA,
                        RScript.ReturnType.INTEGER, Lists.newArrayList(key), seconds);
        if (count > maxCount) {
            long expire = (long) Math.ceil(redissonClient.getBucket(key).remainTimeToLive() / 1000);
            throw new RequestLimiterException(String.format("%s秒内请求次数超过%s次限制!key:%s", seconds, maxCount, key), expire);
        }
    }

    /**
     * 递增请求次数,超过最大允许次数会抛出RequestLimiterException 异常
     *
     * @param key      请求类型KEY
     * @param seconds  单位时间(s)
     * @param maxCount 单位时间内上限次数
     */
    public void incrRequestLimiter(String key, long seconds, long maxCount,Integer code) {
        Long count = redissonClient.getScript()
                .eval(RScript.Mode.READ_WRITE, SCRIPT_LUA,
                        RScript.ReturnType.INTEGER, Lists.newArrayList(key), seconds);
        if (count > maxCount) {
            long expire = (long) Math.ceil(redissonClient.getBucket(key).remainTimeToLive() / 1000);
            throw new RequestLimiterException(code,String.format("%s秒内请求次数超过%s次限制!key:%s", seconds, maxCount, key), expire > 0 ? expire : 1);
        }
    }

    /**
     * 对要求key键增加计数和限时
     *
     * @param key     键
     * @param seconds 秒
     */
    public void incrKeyLimiter(String key, long seconds) {
        redissonClient.getScript().eval(RScript.Mode.READ_WRITE, SCRIPT_LUA, RScript.ReturnType.INTEGER, Lists.newArrayList(key), seconds);
    }

    /**
     * 获取次数
     *
     * @param key 键
     * @return 返回次数
     */
    public Integer getKeyLimiter(String key) {
        Integer count = (Integer) redissonClient.getBucket(key).get();
        if (null != count) {
            return count;
        }
        return 0;
    }

    /**
     * 对要求key键进行计数验证和时间验证
     *
     * @param key      键
     * @param maxCount 最大次数
     * @return 验证失败，返回false
     */
    public boolean checkKeyLimiter(String key, long maxCount) {
        if(maxCount<=0){
            log.info("key[{}]超过最大次数限制{}", key, maxCount);
            return false;
        }
        Integer count = (Integer) redissonClient.getBucket(key).get();
        if (null != count && count >= maxCount) {
            long expire = (long) Math.ceil(redissonClient.getBucket(key).remainTimeToLive() / 1000);
            log.info("key[{}]在失效时间{}前累计次数{}超过最大次数限制{}", key, expire, count, maxCount);
            return false;
        }
        return true;
    }

    /**
     * 请求次数重置
     * 后续事务出现异常, 或者事务成功后
     * 如: 发送验证码出现异常, 或验证码校验成功后; 递减一次;
     *
     * @param identifier  限制标识(客户ID, IP地址,...)
     * @param transaction 事务类型(可以根据需求,传入当前 URI)
     */
    public void resetKeyLimiterForFailure(String transaction, String identifier) {
        final String key = KEY_PREFIX + COLON + transaction + COLON + identifier;
        this.resetKeyLimiter(key);
    }

    /**
     * 对要求key键成功时复位请求次数
     *
     * @param key 键
     */
    public void resetKeyLimiter(String key) {
        redissonClient.getAtomicLong( key).delete();
    }

    /**
     * 递减请求次数
     *
     * @param identifier  限制标识(客户ID, IP地址,...)
     * @param transaction 事务类型(可以根据需求,传入当前 URI)
     * @return long
     */
    public long decrRequestLimiter(String transaction, String identifier) {
        final String key = KEY_PREFIX + COLON + transaction + COLON + identifier;
        return redissonClient.getAtomicLong(key).decrementAndGet();
    }
}
